"""
12_vector_search.py - Vector Databases and Embeddings
======================================================
Chapter 14: The retrieval infrastructure for AI agents.

KEY CONCEPTS:
  ★ Embeddings - Convert text to numerical vectors
  ★ Vector Store - Database for similarity search
  ★ Cosine Similarity - Measure text similarity
  ★ ChromaDB - Lightweight vector database

Run: python src/12_vector_search.py
"""
import math
from typing import List, Tuple

class SimpleVectorStore:
    """Minimal vector store implementation (no dependencies).
    Uses TF-IDF-like vectors for demonstration."""
    
    def __init__(self):
        self.documents = []
        self.vectors = []
        self.vocabulary = set()
    
    def _tokenize(self, text: str) -> list:
        return text.lower().split()
    
    def _text_to_vector(self, text: str) -> dict:
        tokens = self._tokenize(text)
        vector = {}
        for token in tokens:
            vector[token] = vector.get(token, 0) + 1
        return vector
    
    def _cosine_similarity(self, v1: dict, v2: dict) -> float:
        common = set(v1.keys()) & set(v2.keys())
        dot = sum(v1[k] * v2[k] for k in common)
        mag1 = math.sqrt(sum(v**2 for v in v1.values()))
        mag2 = math.sqrt(sum(v**2 for v in v2.values()))
        if mag1 == 0 or mag2 == 0:
            return 0.0
        return dot / (mag1 * mag2)
    
    def add_documents(self, docs: List[str]):
        for doc in docs:
            self.documents.append(doc)
            self.vectors.append(self._text_to_vector(doc))
    
    def search(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
        query_vec = self._text_to_vector(query)
        scores = []
        for i, doc_vec in enumerate(self.vectors):
            sim = self._cosine_similarity(query_vec, doc_vec)
            scores.append((self.documents[i], sim))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]

if __name__ == "__main__":
    print("=" * 60)
    print("12: Vector Search and Embeddings")
    print("=" * 60)
    store = SimpleVectorStore()
    store.add_documents([
        "LangGraph is a framework for building stateful AI agents",
        "Vector databases store embeddings for fast similarity search",
        "RAG combines retrieval with generation for accurate answers",
        "Python is a popular programming language for machine learning",
        "Chatbots use natural language processing to understand users",
        "Multi-agent systems coordinate multiple AI agents together",
    ])
    queries = ["How do AI agents work?", "What is vector search?", "Tell me about Python programming"]
    for q in queries:
        results = store.search(q, top_k=2)
        print(f"\n  Query: {q}")
        for doc, score in results:
            print(f"    [{score:.3f}] {doc}")
    print(f"\n{'=' * 60}\n✅ Vector search complete!")
